
# system pkg >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
import os
import time

from dataclasses import dataclass
from typing import Any, Callable, Iterable

# user pkg >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
import learn_conf as lf

from model import TransformerModel
from model.part import subsequentMask
from tc_datasets import DatasetZhEn
from learn_conf import TensorboardManager as Tbm

# torch pkg >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP

# 创建一个分词器 + 词表
tokenizer = lf.TokenizerTool()

batches_num = 0


@dataclass
# 记录 train 过程中的 step, accum_step, samples, tokens
class Batch:
    src_dat: tuple[torch.Tensor, torch.Tensor]
    tgt_dat: tuple[torch.Tensor, torch.Tensor]
    # src: torch.Tensor        # src shape: (batch_size, src_seq_len)
    # tgt: torch.Tensor        # tgt shape: (batch_size, tgt_seq_len)
    # src_mask: torch.Tensor   # 输入 seq 的 padding mask, shape: (batch_size, src_seq_len)
    # tgt_mask: torch.Tensor   #  shape: (batch_size, tgt_seq_len)


# *****************************************************************************
# brief: train mode / eval mode 下进行一个 epoch 的训练，train 下会自动计算梯度更新参数
# =============================================================================
def run_epoch(
    data_iter: Iterable[Batch],      # 迭代器，每次返回一个 batch（含源、目标、mask 等）
    model: TransformerModel,         # 模型（如 Transformer）
    criterion,
    loss_compute: Callable, # 损失计算函数（通常是一个封装了 loss + backward + scale 的对象）
    optimizer,            # 优化器（如 Adam）
    scheduler,            # 学习率调度器
    mode:str = "train",   # 当前模式：train、eval 或 train+log
    is_main_process: bool = False,  # 是否是 GPU=0 的主进程
    config: dict[str, Any] = None,  # 保存用户信息的字典
    log_writer:lf.TensorboardManager = None,
):
    global batches_num

    start = time.time()  # 当前时间，用于计算速度
    # 本次 epoch 的 总共 tokens
    total_tokens = 0     # 所有 token 的计数（用于最终计算平均 loss）
    # 本次 epoch 的 总共 loss
    total_loss = 0       # 累加的总 loss

    accum_iter = config["accum_iter"]

    device = model.device

    for i, batch in enumerate(data_iter):
        # 处理 batch 数据 >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        # (batch_size, seq_len)
        (en, en_mask), (ch, ch_mask) = batch.src_dat, batch.tgt_dat

        en = en.to(device)
        en_mask = en_mask.to(device)
        ch = ch.to(device)
        # ch_mask = ch_mask.to(device)

        src_seq_len = en.size(-1)
        tgt_seq_len = ch.size(-1) - 1

        # decoder 的输入 (batch_size, tgt_seq_len - 1)
        decoder_input = ch[:, :-1].to(device) # 去掉 最后一个 token
        # decoder 的目标输出 (batch_size, tgt_seq_len - 1)
        tgt_data = ch[:, 1:].to(device)  # 去掉第一个 token "<bos>"

        # decoder 中 self-attention 的 mask : padding mask + subsequent mask
        tgt_mask = (decoder_input != tokenizer.pad_id()).unsqueeze(1).expand(-1, tgt_seq_len, -1)  # padding mask
        tgt_mask = tgt_mask & subsequentMask(tgt_seq_len, tgt_seq_len).to(device) # subsequent mask
        # decoder 中 cross-attention 的 mask : padding mask
        cross_mask = en_mask.bool().unsqueeze(1).expand(-1, tgt_seq_len, -1).to(device)
        # encoder 中 self-attention 的 mask : padding mask
        src_mask = en_mask.bool().unsqueeze(1).expand(-1, src_seq_len, -1).to(device)

        # 总 tokens 数量
        ntokens = (tgt_data != tokenizer.pad_id()).data.sum()

        # model forward >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        out = model(
            src=en, tgt=decoder_input,
            src_mask=src_mask, tgt_mask=tgt_mask, cross_mask=cross_mask
        )

        # loss 是总的 loss 值（标量）, loss_node 是原始的 PyTorch Tensor，可以用来做 .backward()
        loss, loss_node = loss_compute(out, tgt_data, criterion, ntokens)
        # 梯度累积求 mean 的操作
        loss_node = loss_node / accum_iter

        if mode == "train":
            # 反向传播计算梯度，并且计算完之后汇总梯度，保证每个 GPU 的模型上的梯度都相同（同步梯度）
            loss_node.backward()
            # 每隔 accum_iter 个 mini-batch 才调用 optimizer.step() 执行参数更新
            if (i + 1) % accum_iter == 0:
                # 更新 网络中的参数，因为每个 GPU 梯度是一样的，所以更新后的参数也是一样的
                optimizer.step()
                # 清除梯度
                optimizer.zero_grad(set_to_none=True)
                # 更新 学习率（每张 GPU 上面都有一个独立的 scheduler，因此不冲突）
                scheduler.step()

                # 整合所有 GPU 上面的 Loss 数据
                # local_loss = loss_node.detach().to(device)
                # local_ntokens = ntokens.clone().detach().to(device).requires_grad_(False)
                #
                # dist.all_reduce(local_loss, op=dist.ReduceOp.SUM)
                # dist.all_reduce(local_ntokens, op=dist.ReduceOp.SUM)

                batches_num += 1

                # 只需要一个 process 来记录过程 & 保存临时模型参数
                if is_main_process:

                    # print(f"batch_num: {batches_num}")

                    log_writer.write_lr(batches_num, optimizer.param_groups[0]['lr'])
                    # log_writer.write_train_loss(
                    #     batches_num,
                    #     local_loss.item() / local_ntokens.item(),
                    # )
                    log_writer.write_train_loss(batches_num, loss_node.detach().item())
                    log_writer.write_batch(batches_num)

                    if (config is not None and batches_num % config["pt_save_period"] == 0):
                        save_path = os.path.join(config["pt_save_path"], f"tmp_{batches_num}.pt")
                        torch.save({
                            'batches_num': batches_num,
                            'model_state_dict': model.state_dict(),
                            'optimizer_state_dict': optimizer.state_dict(),
                            'scheduler_state_dict': scheduler.state_dict(),
                        }, save_path)

        elif mode == "eval":
            if is_main_process:
                log_writer.write_eval_loss(batches_num, loss_node.detach().item())
        # end of 'if train'

        total_loss += loss
        total_tokens += ntokens

        del loss
        del loss_node
    # end of 'for batch'

    # 计算这个 epoch 总共用了多少时间
    elapsed = time.time() - start
    time_str = time.strftime("%H:%M:%S", time.gmtime(elapsed))

    avg_loss = total_loss / total_tokens
    print(f"[{mode}] Epoch completed in {time_str} | "
          f"Total Loss: {total_loss:.4f} | Total Tokens: {total_tokens} | "
          f"Avg Loss: {avg_loss:.6f}")
    return avg_loss


# *****************************************************************************
# brief: 只是一个 空壳的 optimizer，只是用来 eval 的时候占位的
# =============================================================================
class DummyOptimizer:
    def step(self, closure=None):
        pass

    def zero_grad(self, set_to_none=False):
        pass


# *****************************************************************************
# brief: 只是一个 空壳的 scheduler，只是用来 eval 的时候占位的
# =============================================================================
class DummyScheduler:
    def step(self):
        pass


# *****************************************************************************
# brief: 分布式训练的时候每个 GPU 执行的函数
# =============================================================================
def train_worker(
    gpu,                    # 当前进程负责的 GPU 编号
    ngpus_per_node,         # 当前节点上的 GPU 总数（用于分布式）
    config,                 # 配置字典（包含超参数）
    is_distributed=True,    # 是否进行分布式训练
):
    print(f"Train worker process using GPU: {gpu} for training", flush=True)
    torch.cuda.set_device(gpu)  # 将当前进程绑定到指定 GPU

    # 词典的大小
    vocab_len = tokenizer.vocab_size()

    # [build model in GPUs] >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    # 创建模型并把 model 转移到 GPU 上面，多 GPU 训练的时候，每个 GPU 都有一份 自己的模型
    model = TransformerModel(
        vocab_size=vocab_len,         # 词典大小
        N=config["N"],                # encoder 和 decoder 中 有多少 single layer
        d_model=config["d_model"],    # 做词嵌入的时候，一个 token 编码成的向量的长度
        headsNum=config["heads"],     # 注意力机制中有多少个 注意力头
        d_ff=config["d_ff"],          # encoder, decoder 中 feed-forward 的隐藏层神经元个数
        dropout=config["dropout"],    # 所有 的 dropout 层 dropout 的 概率
    ).cuda(gpu)
    module = model
    # [build model in GPUs] <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    is_main_process = False
    if is_distributed:
        # 初始化通信组
        # "nccl"：使用高性能的 NCCL后端（推荐用于GPU）。
        # init_method = "env://"：通过环境变量传递 rank和 worldsize（如RANK, WORLD_SIZE）。
        # rank = gpu：每个进程的唯一标识，一般与gpu id 相等。
        # world_size = ngpus_per_node：总共的训练进程数量，一般等于GPU个数。
        dist.init_process_group(
            "nccl", init_method="env://", rank=gpu, world_size=ngpus_per_node
        )
        # 将模型包装成 DistributedDataParallel 模型，用于多 GPU 加速训练和梯度同步
        model = DDP(model, device_ids=[gpu])
        module = model.module
        # GPU 代号为 0 的 GPU 是 main process, main process 负责保存权重
        is_main_process = gpu == 0

    log_writer = Tbm("../.run_tensorboard") if is_main_process else None

    # [create dataset] >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    # 创建一个 data manager
    data_zh_en = DatasetZhEn(tokenizer)
    # 在多 GPU 训练（分布式训练）时，通常会把总的 batch size 按 GPU 数量平均分配
    per_gpu_batch_size = config["batch_size"] // ngpus_per_node
    # 创建要训练的 dataloader
    train_dataloader, valid_dataloader = data_zh_en.create_dataloaders(
        batch_size=per_gpu_batch_size,
        world_size=ngpus_per_node,  # 总GPU数
        rank=dist.get_rank(),       # 当前GPU编号
        use_distributed=True
    )
    # [create dataset] <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    optimizer, lr_scheduler = lf.newWarmupOptim(
        model, config["d_model"], config["base_lr"], config["warmup"]
    )

    criterion = lf.LabelSmoothing(
        size=vocab_len, padding_idx=tokenizer.pad_id(), smoothing=0.1
    ).cuda(gpu)

    # [training] >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    for epoch in range(config["num_epochs"]):
        if is_distributed:
            # 进行数据的划分，保证每个 GPU 上都拿到一部分数据，并且进行打乱，
            # 保证 每个 epoch 训练的数据都不同
            train_dataloader.sampler.set_epoch(epoch)

        model.train()

        if is_main_process:
            log_writer.write_epoch(epoch)

        # TRAINING START --------------------------------------------
        print(f"[GPU{gpu}] Epoch {epoch} Training ......", flush=True)
        run_epoch(
            # 生成器表达式（generator expression）, 懒加载的可迭代对象，
            # 不会立即生成所有元素，而是每次迭代时动态计算下一个值
            data_iter=(Batch(en, ch) for en, ch in train_dataloader),
            model=model,
            criterion=criterion,
            loss_compute=lf.computeLoss,
            optimizer=optimizer,
            scheduler=lr_scheduler,
            mode="train",
            is_main_process=is_main_process,
            config=config,
            log_writer=log_writer,
        )
        # 只有主进程（rank==0）负责保存模型，避免多进程重复保存
        if is_main_process:
            file_path = "%s%.2d.pt" % (config["file_prefix"], epoch)
            torch.save(module.state_dict(), file_path)

        # 释放没用的显存，减少 memory leak
        torch.cuda.empty_cache()
        # TRAINING END ----------------------------------------------

        # VALIDATION START ------------------------------------------
        print(f"[GPU{gpu}] Epoch {epoch} Validation ====", flush=True)
        model.eval()
        run_epoch(
            # 生成器表达式（generator expression）, 懒加载的可迭代对象，
            # 不会立即生成所有元素，而是每次迭代时动态计算下一个值
            data_iter=(Batch(en, ch) for en, ch in valid_dataloader),
            model=model,
            criterion=criterion,
            loss_compute=lf.computeLoss,
            optimizer=DummyOptimizer(),
            scheduler=DummyScheduler(),
            mode="eval",
            config=config,
            log_writer=log_writer,
        )

        # 释放没用的显存，减少 memory leak
        torch.cuda.empty_cache()
        # VALIDATION END --------------------------------------------
    # end of 'for epoch'

    if is_main_process:
        log_writer.close()

    # 保存最优权重
    if is_main_process:
        file_path = "%sfinal.pt" % config["file_prefix"]
        torch.save(module.state_dict(), file_path)


# *****************************************************************************
# brief: 进行 GPU 分布式 训练
# =============================================================================
def train_distributed_model(config:dict[str, Any]):
    ngpus = torch.cuda.device_count()
    # 单机多卡训练，通常是 "localhost" 或 "127.0.0.1"，表示本机
    os.environ["MASTER_ADDR"] = "localhost"
    # 指定一个端口来进行进程间通信，端口号需要是一个空闲且没有被占用的端口。
    os.environ["MASTER_PORT"] = "12356"
    print(f"Number of GPUs detected: {ngpus}")
    print("Spawning training processes ...")
    mp.spawn(
        train_worker,
        nprocs=ngpus,
        args=(ngpus, config, True),
    )


if __name__ == "__main__":
    train_distributed_model({
    # 超参数 调整
        "num_epochs":  3,    # 训练多少个 epoch
        "batch_size":  3,     # batch size 是多少
        "base_lr":     0.5,  # 基础学习率，乘上 lr_scheduler 给的系数就是最终学习率
        "warmup":      (300),   # 前多少个 batch 进行 lr 的 warmup
        "accum_iter":  32,      # 每个多少个 batch 更新参数
        "file_prefix": "multi30k_model_",  # pt 文件的前缀
    # 模型参数
        "N": 6, # encoder, decoder 有多少 layer
        "heads": 8, # 多少个注意力头
        "d_model": 512,  # 词嵌入的时候一个 token 转为多少维度的向量
        "d_ff": 1024,  # feed-forward 隐藏层的大小
        "dropout": 0.1, # dropout 的概率

    # 临时 pt 文件的保存
        # 每隔 500 个等效 batch(real batch * accum_iter) 保存一次临时模型
        "pt_save_period": 1000,
        # 临时 pt 文件 的保存路径 (tmp_{equivalent batc num}.pt)
        "pt_save_path": os.path.dirname(os.path.abspath(__file__)),
    })
